# -*- coding: utf-8 -*-
###########################################
# create: 2024-09-25
# author: gql
# desc: hdfs get下载失败的分区进行二次下载修复
############################################
import io
import logging
import os
import sys
import time

from apps.utils.common import CommonUtils
from apps.utils.configparser import ConfigParser
from apps.utils.local import LocalUtils
from apps.utils.logger import Logger
from apps.utils.remote import RemoteTool

reload(sys)
sys.setdefaultencoding('utf8')

config = ConfigParser('../conf/config-get.ini')
config.parse()
# [logger]
log_path = config.get('logger', 'log_path')
log_name = config.get('logger', 'log_name')
# [local]
local_disk_dir = config.get('local', 'local_disk_dir')
local_data_dir = config.get('local', 'local_data_dir')
local_script_dir = config.get('local', 'local_script_dir')
local_part_table_dir = config.get('local', 'local_part_table_dir')
local_big_part_path = config.get('local', 'local_big_part_path')
local_skip_big_path = config.get('local', 'local_skip_big_path')
local_skip_loop_path = config.get('local', 'local_skip_loop_path')
local_disk_limit_percent = config.get('local', 'local_disk_limit_percent')
# [old_cluster]
cls_full_limit = config.get('old_cluster', 'cls_full_limit')
cls_db_name = config.get('old_cluster', 'cls_db_name')
cls_hdfs_auth = config.get('old_cluster', 'cls_hdfs_auth')
# [new_cluster]
remote_host = config.get('new_cluster', 'remote_host')
remote_pass = config.get('new_cluster', 'remote_pass')
remote_port = config.get('new_cluster', 'remote_port')
remote_disk_dir = config.get('new_cluster', 'remote_disk_dir')
remote_data_dir = config.get('new_cluster', 'remote_data_dir')
remote_disk_limit_percent = config.get('new_cluster', 'remote_disk_limit_percent')
remote_record_dir = config.get('new_cluster', 'remote_record_dir')
remote_scp_fin_path = config.get('new_cluster', 'remote_scp_fin_path')
remote_disk_check_secs = config.get('new_cluster', 'remote_disk_check_secs')
# [error]
error_loop_limit = config.get('error', 'error_loop_limit')
error_wait_sec = config.get('error', 'error_wait_sec')
# 日志
log = Logger(log_path, log_name, is_mask=False, c_level=logging.INFO, f_level=logging.INFO)
log.set_passwd(remote_pass)
util = CommonUtils(log)
net = RemoteTool(log, remote_host, remote_pass, remote_port)
net.set_config(remote_disk_dir, remote_disk_limit_percent, remote_disk_check_secs)
loc = LocalUtils(log, util)
loc.set_config(local_disk_dir, local_disk_limit_percent)


# 解析清单文件
def get_manifest_config(manifest_file):
    conf = []
    with io.open(manifest_file, 'r', encoding='utf-8') as fr:
        for line in fr:
            fields = line.strip().split("\t")
            # t_stu	par_dt=20221223	/user/hive/t_stu/par_dt=20221223
            tab_name = fields[0]
            part_name = fields[1]
            abs_path = fields[2]
            is_part = False if part_name == "-" else True
            conf.append((tab_name, part_name, abs_path, is_part))
    return conf


# hdfs下载文件
def hdfs_get(hdfs_path, local_path):
    script = cls_hdfs_auth + "; hdfs dfs -get {} {}".format(hdfs_path, local_path)
    stat_code = util.exec_local_shell0(script)
    return stat_code


# scp完成并校验后，发送信号到远端服务器记录文件
def remote_scp_trans_done(db_name, tab_name, part_name, part_kb):
    script = '"echo {}:{}:{}:{} >> {}"'.format(db_name, tab_name, part_name, part_kb, remote_scp_fin_path)
    net.exec_remote_shell0(script)
    log.info("scp完成信号已传输到远端服务器")


# 下载并传输分区数据
def trans_by_part(tab_name, part_name, part_path, file_skip_loop_tab):
    log.info("hdfs get路径: {}".format(part_path))
    part_local_parent_path = "{}/{}".format(local_data_dir, tab_name)
    part_local_path = "{}/{}/{}".format(local_data_dir, tab_name, part_name)
    stat_code = hdfs_get(part_path, part_local_parent_path)
    loop = 0
    while stat_code is not None and str(stat_code) != "0":
        loop += 1
        if loop == error_loop_limit:
            log.info("下载分区失败{}次,记录并跳过".format(error_loop_limit))
            # 表名 分区 路径
            file_skip_loop_tab.write(u"{}\t{}\t{}\n".format(tab_name, part_name, part_path))
            file_skip_loop_tab.flush()
            loc.remove_path(part_local_path)
            return
        util.sleep_strategy(loop, sleep_sec=error_wait_sec, message="hdfs下载异常(分区)")
        loc.remove_path(part_local_path)
        stat_code = hdfs_get(part_path, part_local_parent_path)

    log.info("hdfs get成功: {}".format(part_path))
    local_part_kb = loc.file_size_kb(part_local_path)
    net.remote_check_disk(local_part_kb)

    log.info("scp远端传输文件(分区)")
    remote_addr = "{}:{}/{}".format(remote_host, remote_data_dir, tab_name)
    net.scp_to_remote(part_local_path, remote_addr)

    # 判断scp完成后目录大小与本地目录大小，如果一致，则写入信息，否则scp重试
    remote_part_path = "{}/{}/{}".format(remote_data_dir, tab_name, part_name)
    remote_part_kb = net.remote_file_size(remote_part_path)
    loop = 0
    while remote_part_kb != local_part_kb:
        loop += 1
        log.error("scp失败, 表名: {}, 分区: {}, Remote_kb: {}, Local_kb: {}".format(
            tab_name, part_name, remote_part_kb, local_part_kb))
        util.sleep_strategy(loop, sleep_sec=60, message="scp失败重传(分区)")
        net.scp_to_remote(part_local_path, remote_addr)
        remote_part_kb = net.remote_file_size(remote_part_path)

    log.info("scp完成远端传输文件并校验成功")
    # 传输完成信息
    remote_scp_trans_done(cls_db_name, tab_name, part_name, local_part_kb)
    loc.remove_path(part_local_path)


# 下载并传输全量数据
def trans_by_full(tab_name, tab_path, file_skip_loop_tab):
    log.info("hdfs get路径: {}".format(tab_path))
    tab_local_parent_path = local_data_dir
    tab_local_path = "{}/{}".format(local_data_dir, tab_name)
    stat_code = hdfs_get(tab_path, tab_local_parent_path)
    loop = 0
    while stat_code is not None and str(stat_code) != "0":
        loop += 1
        if loop == error_loop_limit:
            log.info("下载全量表失败{}次，记录并跳过".format(error_loop_limit))
            # 表名 分区 路径
            file_skip_loop_tab.write(u"{}\t{}\t{}\n".format(tab_name, "-", tab_path))
            file_skip_loop_tab.flush()
            loc.remove_path(tab_local_path)
            return
        util.sleep_strategy(loop, sleep_sec=error_wait_sec, message="hdfs下载异常(全量表)")
        loc.remove_path(tab_local_path)
        stat_code = hdfs_get(tab_path, tab_local_parent_path)

    log.info("hdfs get成功: {}".format(tab_path))
    # 获取本地表大小
    local_tab_kb = loc.file_size_kb(tab_local_path)
    net.remote_check_disk(local_tab_kb)

    log.info("scp远端传输文件(全量表)")
    remote_addr = "{}:{}/".format(remote_host, remote_data_dir)
    net.scp_to_remote(tab_local_path, remote_addr)
    # 判断scp完成后目录大小与本地目录大小，如果一致，则写入信息，否则重试scp
    remote_tab_path = "{}/{}".format(remote_data_dir, tab_name)
    remote_tab_kb = net.remote_file_size(remote_tab_path)
    loop = 0
    while remote_tab_kb != local_tab_kb:
        loop += 1
        log.error("scp失败，表名: {}, Remote_kb: {}, Local_kb: {}".format(tab_name, remote_tab_kb, local_tab_kb))
        # 睡眠等待，scp失败重传
        util.sleep_strategy(loop, sleep_sec=60, message="scp失败重传(全量表)")
        net.scp_to_remote(tab_local_path, remote_addr)
        remote_tab_kb = net.remote_file_size(remote_tab_path)

    log.info("scp远端传输文件完成并校验成功")
    # 传输完成信息
    remote_scp_trans_done(cls_db_name, tab_name, "", local_tab_kb)
    loc.remove_path(tab_local_path)


def main(argv):
    if len(argv) <= 2:
        print("""请输入2个参数: 批次号 批次清单文件
        比如: python {} 001 001.txt""".format(argv[0]))
        exit(2)
    total_kb, used_kb, avail_kb, use_percent = loc.local_disk_size()
    if float(use_percent.replace("%", "")) > float(local_disk_limit_percent) * 100:
        log.warn("磁盘使用率超过 {}%，请核对！".format(local_disk_limit_percent))
        exit(3)
    # 批次号
    batch_no = argv[1]
    # 清单文件
    manifest_file = argv[2]
    start_dt = util.date_format('%Y-%m-%d %H:%M:%S')
    start_time = time.time()

    net.exec_remote_shell0("mkdir -p {}".format(remote_data_dir))
    net.exec_remote_shell0("mkdir -p {}".format(remote_record_dir))
    file_skip_loop_tab = io.open(local_skip_loop_path, 'a+', encoding='utf-8')

    log.info("解析清单文件: {}".format(manifest_file))
    config_list = get_manifest_config(manifest_file)
    for tab_name, part_name, abs_path, is_part in config_list:
        log.info("---> 正在执行：{} {} {} {}".format(tab_name, part_name, abs_path, is_part))
        table_local_path = "{}/{}".format(local_data_dir, tab_name)
        if not os.path.exists(table_local_path):
            os.makedirs(table_local_path)
        remote_tab_path = remote_data_dir + "/" + tab_name
        script = '"mkdir -p {}"'.format(remote_tab_path)
        net.exec_remote_shell0(script)
        if is_part:
            trans_by_part(tab_name, part_name, abs_path, file_skip_loop_tab)
        else:
            trans_by_full(tab_name, abs_path, file_skip_loop_tab)
    file_skip_loop_tab.close()

    end_time = time.time()
    log.info("*" * 50)
    log.info("\t完成批次：{}".format(batch_no))
    log.info("\t开始时间：" + start_dt)
    log.info("\t结束时间：" + util.date_format('%Y-%m-%d %H:%M:%S'))
    log.info("\t耗费时间：{:.2f}秒".format(end_time - start_time))
    log.info("*" * 50)


if __name__ == '__main__':
    try:
        main(sys.argv)
    except Exception as e:
        log.error("出现错误主程序退出! {}".format(e))
